import os
import zipfile
import tarfile
import py7zr
from typing import List, Optional
import pandas as pd

class ArchiveScanner:
    """
    通用压缩包内容扫描器
    支持格式: ZIP, RAR, TAR, 7Z, TAR.GZ, TAR.BZ2 等
    """
    
    def __init__(self):
        self.supported_formats = ['.zip', '.rar', '.tar', '.7z', '.gz', '.bz2', '.tgz', '.tbz2']
    
    def get_archive_contents(self, archive_path: str, output_file: Optional[str] = None, save2xlsx: bool = False) -> List[str]:
        """
        获取压缩包内所有文件的完整路径
        
        Args:
            archive_path: 压缩包文件路径
            output_file: 可选输出文件路径
            save2xlsx: 是否保存为Excel文件(.xlsx格式)
            
        Returns:
            包含所有文件路径的列表
        """
        if not os.path.exists(archive_path):
            raise FileNotFoundError(f"压缩包文件不存在: {archive_path}")
        
        file_ext = os.path.splitext(archive_path)[1].lower()
        all_files = []
        
        try:
            # ZIP 格式
            if file_ext == '.zip':
                all_files = self._scan_zip(archive_path)
            
            # TAR 及相关格式
            elif file_ext in ['.tar', '.gz', '.bz2', '.tgz', '.tbz2']:
                all_files = self._scan_tar(archive_path)
            
            # 7Z 格式
            elif file_ext == '.7z':
                all_files = self._scan_7z(archive_path)
            
            # RAR 格式
            elif file_ext == '.rar':
                all_files = self._scan_rar(archive_path)
            
            else:
                raise ValueError(f"不支持的压缩格式: {file_ext}")
            
            # 写入输出文件
            if output_file and all_files:
                if save2xlsx:
                    self._write_to_excel(all_files, output_file)
                else:
                    self._write_to_file(all_files, output_file)
                
            return all_files
            
        except Exception as e:
            # 检查是否是密码保护相关的错误
            error_msg = str(e).lower()
            if "password" in error_msg or "加密" in error_msg or "encrypted" in error_msg:
                raise Exception(f"压缩包 {archive_path} 受密码保护，无法读取内容。请提供密码或使用无密码的压缩包。")
            elif "bad" in error_msg or "invalid" in error_msg:
                raise Exception(f"压缩包 {archive_path} 损坏或格式无效，无法读取内容。")
            else:
                raise Exception(f"处理压缩包 {archive_path} 时出错: {str(e)}")
    
    def scan_directory_recursively(self, directory_path: str, recursive: bool = True) -> List[str]:
        """
        递归扫描目录中的所有文件和压缩包
        
        Args:
            directory_path: 目录路径
            recursive: 是否递归扫描子目录
            
        Returns:
            包含所有文件路径的列表
        """
        if not os.path.exists(directory_path):
            raise FileNotFoundError(f"目录不存在: {directory_path}")
        
        if not os.path.isdir(directory_path):
            raise ValueError(f"路径不是目录: {directory_path}")
        
        all_files = []
        
        try:
            if recursive:
                # 递归扫描目录
                for root, dirs, files in os.walk(directory_path):
                    for file in files:
                        file_path = os.path.join(root, file)
                        relative_path = os.path.relpath(file_path, directory_path)
                        all_files.append(relative_path)
            else:
                # 只扫描顶层目录
                for item in os.listdir(directory_path):
                    item_path = os.path.join(directory_path, item)
                    if os.path.isfile(item_path):
                        all_files.append(item)
                        
        except Exception as e:
            raise Exception(f"扫描目录 {directory_path} 时出错: {str(e)}")
        
        return sorted(all_files)
    
    def get_recursive_contents(self, archive_path: str, output_file: Optional[str] = None, save2xlsx: bool = False) -> List[str]:
        """
        递归获取压缩包内所有文件的完整路径（包括嵌套的压缩包）
        
        Args:
            archive_path: 压缩包文件路径
            output_file: 可选输出文件路径
            save2xlsx: 是否保存为Excel文件(.xlsx格式)
            
        Returns:
            包含所有文件路径的列表
        """
        if not os.path.exists(archive_path):
            raise FileNotFoundError(f"压缩包文件不存在: {archive_path}")
        
        # 获取当前压缩包的所有内容
        all_files = self.get_archive_contents(archive_path)
        
        # 递归处理嵌套的压缩包
        recursive_files = []
        for file_path in all_files:
            # 检查是否为支持的压缩包格式
            file_ext = os.path.splitext(file_path)[1].lower()
            if file_ext in self.supported_formats:
                # 注意：这里我们只能获取嵌套压缩包内的文件列表，
                # 但无法直接访问嵌套压缩包的内容，除非先解压
                # 这里只是标记嵌套压缩包的存在
                recursive_files.append(f"[嵌套压缩包] {file_path}")
            else:
                recursive_files.append(file_path)
        
        # 写入输出文件
        if output_file and recursive_files:
            if save2xlsx:
                self._write_to_excel(recursive_files, output_file)
            else:
                self._write_to_file(recursive_files, output_file)
            
        return recursive_files
    
    def process_directory_archives(self, directory_path: str, recursive: bool = True, save2xlsx: bool = False) -> dict:
        """
        处理目录中的所有压缩包文件
        
        Args:
            directory_path: 目录路径
            recursive: 是否递归处理嵌套压缩包
            save2xlsx: 是否保存为Excel文件(.xlsx格式)
            
        Returns:
            字典，键为压缩包文件名，值为文件列表
        """
        if not os.path.exists(directory_path):
            raise FileNotFoundError(f"目录不存在: {directory_path}")
        
        results = {}
        
        # 获取目录中的所有文件
        files = self.scan_directory_recursively(directory_path, False)  # 只扫描顶层目录
        
        for file in files:
            file_path = os.path.join(directory_path, file)
            file_ext = os.path.splitext(file_path)[1].lower()
            
            # 检查是否为支持的压缩包格式
            if file_ext in self.supported_formats:
                try:
                    print(f"正在处理压缩包: {file}")
                    if recursive:
                        contents = self.get_recursive_contents(file_path, save2xlsx=save2xlsx)
                    else:
                        contents = self.get_archive_contents(file_path, save2xlsx=save2xlsx)
                    results[file] = contents
                except Exception as e:
                    print(f"处理 {file} 时出错: {e}")
                    results[file] = [f"错误: {str(e)}"]
        
        return results
    
    def _scan_zip(self, archive_path: str) -> List[str]:
        """扫描 ZIP 格式压缩包"""
        all_files = []
        try:
            with zipfile.ZipFile(archive_path, 'r') as zip_ref:
                for file_info in zip_ref.infolist():
                    # 排除目录，只保留文件
                    if not file_info.is_dir():
                        # 标准化路径（将反斜杠转换为正斜杠）
                        file_path = file_info.filename.replace('\\', '/')
                        # 处理编码问题
                        try:
                            # 尝试使用UTF-8解码
                            decoded_path = file_path.encode('cp437').decode('utf-8')
                            all_files.append(decoded_path)
                        except (UnicodeDecodeError, UnicodeEncodeError):
                            # 如果UTF-8解码失败，尝试其他编码或保持原样
                            try:
                                # 尝试使用GBK解码（常用于中文）
                                decoded_path = file_path.encode('cp437').decode('gbk')
                                all_files.append(decoded_path)
                            except (UnicodeDecodeError, UnicodeEncodeError):
                                # 如果都失败了，保持原样
                                all_files.append(file_path)
        except zipfile.BadZipFile as e:
            raise ValueError(f"无效的ZIP文件: {str(e)}")
        except RuntimeError as e:
            # 检查是否是密码保护相关的错误
            if "password" in str(e).lower():
                raise ValueError(f"ZIP文件受密码保护: {str(e)}")
            else:
                raise ValueError(f"读取ZIP文件时出错: {str(e)}")
        return all_files
    
    def _scan_tar(self, archive_path: str) -> List[str]:
        """扫描 TAR 及相关格式压缩包"""
        all_files = []
        try:
            with tarfile.open(archive_path, 'r:*') as tar_ref:
                for member in tar_ref.getmembers():
                    if member.isfile():
                        # 标准化路径
                        file_path = member.name.replace('\\', '/')
                        all_files.append(file_path)
        except tarfile.ReadError as e:
            raise ValueError(f"无效的TAR文件: {str(e)}")
        return all_files
    
    def _scan_7z(self, archive_path: str) -> List[str]:
        """扫描 7Z 格式压缩包"""
        all_files = []
        try:
            with py7zr.SevenZipFile(archive_path, 'r') as sevenz_ref:
                for file_info in sevenz_ref.list():
                    if not file_info.is_directory:
                        file_path = file_info.filename.replace('\\', '/')
                        all_files.append(file_path)
        except Exception as e:
            # 检查是否是密码保护相关的错误
            error_msg = str(e).lower()
            if "password" in error_msg or "encrypted" in error_msg:
                raise ValueError(f"7Z文件受密码保护: {str(e)}")
            else:
                raise ValueError(f"无效的7Z文件: {str(e)}")
        return all_files
    
    def _scan_rar(self, archive_path: str) -> List[str]:
        """扫描 RAR 格式压缩包"""
        all_files = []
        try:
            # 尝试使用 rarfile 库
            import rarfile
            # 设置RAR文件的编码
            rarfile.UNICODE_COMMENTS = True
            rarfile.PATH_SEP = '/'
            with rarfile.RarFile(archive_path, 'r') as rar_ref:
                for file_info in rar_ref.infolist():
                    if not file_info.isdir():
                        file_path = file_info.filename.replace('\\', '/')
                        all_files.append(file_path)
        except ImportError:
            # 回退到使用 unrar 命令（如果系统安装了 unrar）
            all_files = self._scan_rar_fallback(archive_path)
        except Exception as e:
            # 检查是否是密码保护相关的错误
            error_msg = str(e).lower()
            if "password" in error_msg or "加密" in error_msg:
                raise ValueError(f"RAR文件受密码保护: {str(e)}")
            else:
                raise ValueError(f"无效的RAR文件: {str(e)}")
        return all_files
    
    def _scan_rar_fallback(self, archive_path: str) -> List[str]:
        """使用系统命令作为 RAR 扫描的备选方案"""
        try:
            import subprocess
            result = subprocess.run(['unrar', 'lb', archive_path], 
                                  capture_output=True, text=True, check=True, encoding='utf-8')
            files = [line.strip() for line in result.stdout.splitlines() if line.strip()]
            return [f.replace('\\', '/') for f in files]
        except subprocess.CalledProcessError as e:
            # 检查是否是密码保护相关的错误
            error_output = e.stderr.lower()
            if "password" in error_output or "加密" in error_output:
                raise ImportError(f"RAR文件受密码保护，请安装 rarfile 库: pip install rarfile 或提供密码")
            else:
                raise ImportError(f"请安装 rarfile 库: pip install rarfile 或系统安装 unrar")
        except FileNotFoundError:
            raise ImportError("请安装 rarfile 库: pip install rarfile 或系统安装 unrar")
    
    def _write_to_file(self, file_list: List[str], output_file: str):
        """将文件列表写入到输出文件"""
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                for file_path in sorted(file_list):
                    f.write(file_path + '\n')
            print(f"文件列表已保存到: {output_file}")
        except Exception as e:
            print(f"写入输出文件失败: {str(e)}")
    
    def _write_to_excel(self, file_list: List[str], output_file: str):
        """将文件列表保存到Excel文件"""
        try:
            # 确保输出文件名以.xlsx结尾
            if not output_file.endswith('.xlsx'):
                output_file = output_file.rsplit('.', 1)[0] + '.xlsx'
            
            # 创建DataFrame，只包含一列相对路径
            df = pd.DataFrame({'相对路径': sorted(file_list)})
            
            # 保存到Excel文件
            df.to_excel(output_file, index=False, engine='openpyxl')
            print(f"文件列表已保存到Excel文件: {output_file}")
        except Exception as e:
            print(f"写入Excel文件失败: {str(e)}")
    
    def is_supported_format(self, file_path: str) -> bool:
        """检查文件格式是否支持"""
        file_ext = os.path.splitext(file_path)[1].lower()
        return file_ext in self.supported_formats


# 使用示例
def main():
    scanner = ArchiveScanner()
    
    # 设置查找路径
    target_path = r"input_files/jetbra.zip"
    
    # 检查路径是文件还是目录
    if os.path.exists(target_path):
        if os.path.isfile(target_path):
            # 如果是文件，按压缩包处理
            print(f"处理压缩包文件: {target_path}")
            try:
                # 使用递归参数控制是否递归处理嵌套压缩包
                recursive_processing = True
                if recursive_processing:
                    files = scanner.get_recursive_contents(target_path, f"{target_path}_contents_recursive.txt", save2xlsx=True)
                else:
                    files = scanner.get_archive_contents(target_path, f"{target_path}_contents.txt", save2xlsx=True)
                
                print(f"找到 {len(files)} 个文件:")
                for i, file_path in enumerate(files[:20], 1):  # 只显示前20个
                    print(f"  {i:2d}. {file_path}")
                if len(files) > 20:
                    print(f"  ... 还有 {len(files) - 20} 个文件")
            except Exception as e:
                print(f"处理压缩包时出错: {e}")
        elif os.path.isdir(target_path):
            # 如果是目录，按目录处理
            print(f"扫描目录: {target_path}")
            try:
                # 使用递归参数控制是否递归扫描子目录
                recursive_scanning = True
                if recursive_scanning:
                    # 递归处理目录中的所有压缩包
                    archive_results = scanner.process_directory_archives(target_path, recursive=True, save2xlsx=True)
                    
                    # 输出结果
                    print(f"目录中找到 {len(archive_results)} 个压缩包:")
                    total_files = 0
                    for archive_name, contents in archive_results.items():
                        print(f"\n{archive_name}: {len(contents)} 个项目")
                        total_files += len(contents)
                        for i, file_path in enumerate(contents[:10], 1):  # 只显示前10个
                            print(f"  {i:2d}. {file_path}")
                        if len(contents) > 10:
                            print(f"  ... 还有 {len(contents) - 10} 个项目")
                    
                    print(f"\n总计: {total_files} 个项目")
                    
                    # 保存详细结果到文件，使用UTF-8编码
                    detailed_output = "solidworks_macros_detailed_contents_fixed.xlsx"
                    # 创建一个汇总的Excel文件
                    all_data = []
                    for archive_name, contents in archive_results.items():
                        for file_path in contents:
                            all_data.append({'压缩包': archive_name, '相对路径': file_path})
                    
                    df = pd.DataFrame(all_data)
                    df.to_excel(detailed_output, index=False, engine='openpyxl')
                    print(f"\n详细结果已保存到Excel文件: {detailed_output}")
                else:
                    files = scanner.scan_directory_recursively(target_path, recursive=False)
                    
                    print(f"找到 {len(files)} 个文件:")
                    for i, file_path in enumerate(files[:20], 1):  # 只显示前20个
                        print(f"  {i:2d}. {file_path}")
                    if len(files) > 20:
                        print(f"  ... 还有 {len(files) - 20} 个文件")
                        
                    # 保存到Excel文件
                    output_file = f"{os.path.basename(target_path)}_directory_contents.xlsx"
                    scanner._write_to_excel(files, output_file)
            except Exception as e:
                print(f"扫描目录时出错: {e}")
    else:
        print(f"路径不存在: {target_path}")


# 简单使用函数
def get_compressed_contents(archive_path: str, output_file: Optional[str] = None, save2xlsx: bool = False) -> List[str]:
    """
    快速获取压缩包内容的简便函数
    
    Args:
        archive_path: 压缩包路径
        output_file: 可选输出文件路径
        save2xlsx: 是否保存为Excel文件(.xlsx格式)
        
    Returns:
        文件路径列表
    """
    scanner = ArchiveScanner()
    return scanner.get_archive_contents(archive_path, output_file, save2xlsx)


class AdvancedArchiveScanner(ArchiveScanner):
    """增强版压缩包扫描器"""
    
    def get_contents_with_info(self, archive_path: str) -> List[dict]:
        """获取带详细信息的文件列表"""
        file_ext = os.path.splitext(archive_path)[1].lower()
        files_info = []
        
        if file_ext == '.zip':
            with zipfile.ZipFile(archive_path, 'r') as zip_ref:
                for info in zip_ref.infolist():
                    if not info.is_dir():
                        files_info.append({
                            'path': info.filename.replace('\\', '/'),
                            'size': info.file_size,
                            'compressed_size': info.compress_size,
                            'date': f"{info.date_time[0]}-{info.date_time[1]:02d}-{info.date_time[2]:02d}"
                        })
        
        elif file_ext in ['.tar', '.gz', '.bz2']:
            with tarfile.open(archive_path, 'r:*') as tar_ref:
                for member in tar_ref.getmembers():
                    if member.isfile():
                        files_info.append({
                            'path': member.name.replace('\\', '/'),
                            'size': member.size,
                            'date': f"{member.mtime}",
                            'type': member.type
                        })
        
        return files_info
    
    def search_in_archive(self, archive_path: str, pattern: str) -> List[str]:
        """在压缩包中搜索包含特定模式的文件"""
        all_files = self.get_archive_contents(archive_path)
        import fnmatch
        return [f for f in all_files if fnmatch.fnmatch(f.lower(), f'*{pattern.lower()}*')]


if __name__ == "__main__":
    main()